import os

import psutil
from langchain import hub
from langchain.agents import initialize_agent
from langchain_community.llms.tongyi import Tongyi
from langchain_community.tools import TavilySearchResults
from langchain_core.tools import Tool

os.environ["TAVILY_API_KEY"] = "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"
os.environ["DASHSCOPE_API_KEY"] = 'sk-5118bb72818342d1b79539d78c3fbba7'


def get_cpu_info():
    cpu_info = {
        "cpu_count": psutil.cpu_count(),
        "cpu_percent": psutil.cpu_percent(interval=1),
        "cpu_freq": psutil.cpu_freq()._asdict()
    }
    return cpu_info


def get_memory_info():
    memory = psutil.virtual_memory()._asdict()  # 获取虚拟内存信息
    return memory


def get_disk_info():
    disk = psutil.disk_usage('/')._asdict()  # 获取根目录的磁盘使用信息
    return disk


memory_info_tool = Tool(
    name="MemoryInfoTool",
    func=lambda x: get_memory_info(),
    description="用于获取内存的详细信息，包括总内存、已使用、可用内存等"
)

# 定义磁盘信息工具
disk_info_tool = Tool(
    name="DiskInfoTool",
    func=lambda x: get_disk_info(),
    description="用于获取磁盘使用情况，包括总空间、已使用、可用空间等"
)

cpu_info_tool = Tool(
    name="CPUInfoTool",
    func=lambda x: get_cpu_info(),
    description="用于获取当前PC的CPU的详细信息，包括CPU使用率、核心数、频率等"
)

llm = Tongyi()
tools = [TavilySearchResults(), cpu_info_tool, disk_info_tool, memory_info_tool]
prompt = hub.pull("hwchase17/openai-functions-agent")
agent = initialize_agent(tools,
                         llm,
                         agent="zero-shot-react-description",
                         verbose=True)

agent.run("获取当前PC的CPU负载情况？")
agent.run("获取当前PC的内存占用情况？")
agent.run("获取当前PC的磁盘使用情况？")
